#ifndef KLEE_WICFG_H
#define KLEE_WICFG_H
#include <vector>
#include <unordered_map>
#include <set>
#include <memory>
#include <iostream>

#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/Support/raw_os_ostream.h"

#include "klee/Module/KModule.h"
#include "klee/Expr/Constraints.h"
// #include "klee/Core/Executor"

namespace llvm {
    class Function;
    class BasicBlock;
    class Instruction;
    class Module; 
}

namespace klee {
    class Executor;
    class ExecutionState;
    class KInstruction;
    class CallChain_Node;
    class CallChain;
    class WICFG_Node;    
    class WICFG;
    class ConstraintSet;
}

namespace klee {
class CallChain_Node {
    friend class CallChain;
    friend class WICFG;
    using NodePtr = std::shared_ptr<CallChain_Node>;
    using ID_t = unsigned long long;
public:
    CallChain_Node(llvm::Function *_CallFunc, llvm::BasicBlock *_CallBB = nullptr, llvm::Instruction *_CallInst = nullptr)
    : ID(_id++), CallFunc(_CallFunc), CallBB(_CallBB), CallInst(_CallInst), predecessors(nullptr) {}
private:
    static ID_t _id;
    ID_t ID;
    llvm::Function *CallFunc;
    llvm::BasicBlock *CallBB;
    llvm::Instruction *CallInst;
    NodePtr predecessors;
    std::vector<NodePtr> successors;
    std::string getBody();
    std::string getCircleBody();
};

class CallChain {
    friend class WICFG;
    using NodePtr = std::shared_ptr<CallChain_Node>;
    CallChain_Node::NodePtr entryNode;
    std::set<NodePtr> alreadyPrinted; // 避免重复打印.dot文件中的节点
    std::string printSuccessor(NodePtr parent);
    void printToDotFile();
};

// for rapid search of feasible path in loop
struct ConstraintTreeNode { // 保存一个分支判断处的约束
    using NodePtr = std::shared_ptr<ConstraintTreeNode>;
    using ID_t = unsigned long long;
    static ID_t _id;
    ConstraintTreeNode(ref<Expr> _cond = nullptr, ConstraintTreeNode::NodePtr _parent = nullptr, bool _isLoop = false):
        constraint(_cond), parent(_parent), isLoop(_isLoop), ID(_id++) {}
    ref<Expr> constraint; // 当前节点分支对应的约束，如果对应loop节点，则只保存false分支对应的约束
    NodePtr parent; // 当前节点的父节点，如果为空即为根节点
    bool isLoop = false; // 标记当前节点对应WICFG节点是否为loop节点
    ID_t ID;
    std::vector<NodePtr> next; // 当前分支节点对应的直接分支子节点，如果为空即为叶节点
    std::string getBody();
};
struct ConstraintTree { // 前驱WICFG节点中没有Loop节点的Loop节点才可能会有对应的约束树
    using Ptr = std::shared_ptr<ConstraintTree>;
    ConstraintTree(ConstraintTreeNode::NodePtr _root):
        root(_root){}
    ConstraintTreeNode::NodePtr root;
    void removeChild(ConstraintTreeNode::NodePtr &child);
    std::string printSuccessor(ConstraintTreeNode::NodePtr parent);
    void printToDotfile();
};

/// WICFG stands for Weighted Inter-procedure Control Flow Graph
/// View it as a map that guide seletor to find a candidate state which lead to a greatest value
class WICFG_Node {
    using weight_t = long long;
    using NodePtr = std::shared_ptr<WICFG_Node>;
    using ID_t = unsigned long long;
public:
    WICFG_Node() = delete;
    WICFG_Node(WICFG_Node&) = default;
    WICFG_Node(llvm::Function *_func, llvm::Instruction *_callerInst, ID_t _ID, ID_t _callSeries, bool _isTarget = false)
    : func(_func), callerInst(_callerInst), ID(_ID), callSeries(_callSeries), isTarget(_isTarget),
      isUncovered(true), isEncounterRecursion(false), isEncounterLoop(false), numOfInst(0), weight(0) {}
    WICFG_Node& operator=(const WICFG_Node& rhs) = default;
    void setConstraintTree(unsigned id, ConstraintTree::Ptr tree) { constraintTreeMap[id] = tree; }
    friend class WICFG;
    friend class Executor;
    friend class ExecutionState;
    friend class CovTargetOrientedSearcher;
    friend class Searcher;
    // friend void klee::Executor::executeInstruction(ExecutionState &state, KInstruction *ki);
private:
    std::string getBody();
    std::string getCircleBody();
    llvm::Function *func; // 关联函数，为了后续查找调用链，匹配对应的candidate state
    // 关联调用函数中具体发起调用的call指令，只记录调用函数的指针，粒度不够细，不能匹配到关联的候选状态
    llvm::Instruction *callerInst;
    ID_t ID;
    ID_t callSeries;
    std::vector<NodePtr> predecessors;
    std::vector<NodePtr> successors;
    bool isTarget; // 标记该节点对应的函数是否为被测目标函数
    bool isUncovered; // 标记该节点对应的基本块是否已覆盖
    bool isEncounterRecursion = false; // 标记是否遇到循环递归调用
    bool isEncounterLoop = false; // 标记是否遇到循环
    unsigned numOfInst; // 标记该节点对应的指令数量
    weight_t weight; // 标记该节点对应的权重，在有新目标节点被覆盖的时候，需要同步更新。
    unsigned distanceToEnd;
    std::vector<std::string> opcodeNameSeries;
    // 保存由缘起ExecutionState的id和约束树的映射，因为loop节点的前驱路径可能有多种，所以用stateID来区分
    // 保存这一项的主要目的是为了防止loop节点对应的基本块重复创建约束树
    // 约束树保存从本WICFG节点开始，遇到的所有约束，只有当isEncounterLoop为true的时候才可能有内容
    // ConstraintTree::Ptr constraintTree;
    std::map<unsigned, ConstraintTree::Ptr> constraintTreeMap;
};

class WICFG {
    friend class Executor;
    friend class ExecutionState;
    friend class CovTargetOrientedSearcher;
    friend class Searcher;
    using NodePtr = WICFG_Node::NodePtr;
    using weight_t = WICFG_Node::weight_t;
    // using BB2WICFGNode_t = std::unordered_map<llvm::BasicBlock*, NodePtr>;
    using BB2WICFGNodeVec_t = std::unordered_map<llvm::BasicBlock*, std::vector<NodePtr>>;
    using Inst2WICFGNodeVec_t = std::unordered_map<llvm::Instruction*, std::vector<NodePtr>>;
    using FuncSeries_t = unsigned long long;
    using ID_t = unsigned long long;
public:
    WICFG() = delete;
    WICFG(llvm::Function *targetFunc)
    : numOfTargetFuncInst(0), numOfTargetFuncInstCovered(0), numOfCalledTime(0), numOfWICFG_Node(0)
    { init(targetFunc); }
    llvm::Function *getTargetFunc() { return targetFunc; }
    NodePtr getEntryNode() { return entryNode; }
    double getTargetCoveredRatio() { return numOfTargetFuncInstCovered / numOfTargetFuncInst; }
    void fillBB2WICFGNodeVec(BB2WICFGNodeVec_t &mp, llvm::BasicBlock *BB, NodePtr curNode);
    void fillInst2WICFGNodeVec(Inst2WICFGNodeVec_t &mp, llvm::Instruction *inst, NodePtr curNode);
    NodePtr mapBB2WICFGNode(llvm::BasicBlock*);
    void printToDotFile();
    bool isHaltExecution() { return haltExecution; }
    WICFG_Node *getRelatedWICFGNode(llvm::Instruction *inst, llvm::Instruction *callInst);
    static const weight_t initWeight = 10;
private:
    void init(llvm::Function *targetFunc);
    NodePtr traverseFunc(llvm::Function*, llvm::Instruction*, NodePtr, CallChain_Node::NodePtr, bool, FuncSeries_t);
    NodePtr traverseBasicBlock(llvm::BasicBlock*, llvm::Instruction*, NodePtr, NodePtr&, CallChain_Node::NodePtr, bool, FuncSeries_t, bool);
    void changeWeight(NodePtr, weight_t, int, bool, std::set<NodePtr>&);
    void changeWeight2(NodePtr, weight_t, int, bool, std::set<NodePtr>&);
    void setDistance(NodePtr targetNode, int pathLen, std::set<NodePtr>& alreadyChangedNode);
    bool checkIsEncounterLoop(NodePtr curNode, NodePtr startNode, std::set<NodePtr>& inProcess);
    bool checkIsEncounterRecursion(CallChain_Node::NodePtr, llvm::Function*);
    std::string printSuccessor(NodePtr parent);
    std::string printPredecessor(NodePtr child);
    NodePtr entryNode; // 入口节点，不对应实际的基本块，辅助用
    NodePtr endNode; // deprecated
    CallChain callChain;
    std::vector<BB2WICFGNodeVec_t> BB2WICFGNodeMapVec; // 确保同一调用点的函数中的基本块不重复创建WICFG节点
    BB2WICFGNodeVec_t BB2WICFGNodeVecMap; // 为了后续通过指令关联WICFG节点，deprecated
    Inst2WICFGNodeVec_t Inst2WICFGNodeVecMap; // 为了后续通过指令关联WICFG节点
    std::set<NodePtr> alreadyPrinted; // 避免重复打印.dot文件中的节点
    std::set<NodePtr> uncoveredTargetWICFG_Node; // 目标被测函数中未覆盖BasicBlock对应的WICFG节点
    llvm::Function *targetFunc;
    unsigned numOfTargetFuncInst;
    unsigned numOfTargetFuncInstCovered;
    unsigned numOfCalledTime;
    ID_t numOfWICFG_Node;
    bool haltExecution = false;
};
};

#endif